// Copyright lowRISC contributors (OpenTitan project).
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

#include "sw/device/lib/crypto/impl/integrity.h"
#include "sw/device/lib/crypto/include/datatypes.h"
#include "sw/device/lib/crypto/include/hmac.h"
#include "sw/device/lib/crypto/include/sha2.h"
#include "sw/device/lib/runtime/log.h"
#include "sw/device/lib/testing/rand_testutils.h"
#include "sw/device/lib/testing/test_framework/check.h"
#include "sw/device/lib/testing/test_framework/ottf_main.h"

// The autogen rule that creates this header creates it in a directory named
// after the rule, then manipulates the include path in the
// cc_compilation_context to include that directory, so the compiler will find
// the version of this file matching the Bazel rule under test.
#include "hmac_testvectors.h"

// Module ID for status codes.
#define MODULE_ID MAKE_MODULE_ID('t', 's', 't')

// We need the following assertion, because we are using hash context struct
// also for hmac contexts.
static_assert(sizeof(otcrypto_sha2_context_t) ==
                  sizeof(otcrypto_hmac_context_t),
              "Hash and Hmac contexts are expected to be of the same length");

/**
 * This enum defines the different stages of a test vector during streaming
 * or oneshot HMAC calls.
 *
 * Oneshot calls transition from `kHmacTestNotStarted` to `kHmacTestDone`.
 *
 * For streaming calls, the exact number of streaming segments can be up to 4.
 * For fewer segments, a transition from `kHmacTestFeedSegmentX` to
 * `kHmacTestFinalize` is allowed so long as all message bytes of the test is
 * fed to HMAC core.
 *
 * The value of `kHmacTestFeedSegmentX` enum is arranged to match integer X,
 * as it us used as an index to access `segment_idx` array defined in
 * `hmac_extended_test_vector_t`.
 */
typedef enum hmac_test_progress {
  /* The oneshot or streaming test has not been initialized.*/
  kHmacTestNotStarted = 0,
  /* The streaming test is initialized, but no message segment is fed.*/
  kHmacTestFeedSegment1 = 1,
  /* The streaming test is initialized and the first message segment is fed.*/
  kHmacTestFeedSegment2 = 2,
  /* The streaming test is initialized and the first two message segments are
     fed.*/
  kHmacTestFeedSegment3 = 3,
  /* The streaming test is initialized and the first three message segments are
     fed.*/
  kHmacTestFeedSegment4 = 4,
  /* All message bytes are fed but the final is not called.*/
  kHmacTestFinalize = 5,
  /* The test is completed and the result is verified.*/
  kHmacTestDone = 6,
} hmac_test_progress_t;

/**
 * Extend `hmac_test_vector_t` with metadata fields useful for streaming.
 */
typedef struct hmac_extended_test_vector {
  /* The pointer to the test vector coming from the autogenerated C header. */
  hmac_test_vector_t *hmac_test_vector;
  /* The original index of this test vector in the autogenerated C header. */
  size_t original_idx;
  /* The number of update calls for a streaming test. */
  size_t segment_count;
  /* Randomly generated indices to break down the test message to segments. */
  size_t segment_idx[4];
  /* `progess` keeps track of how many message segments are streamed so far. */
  hmac_test_progress_t progress;
  /* `hash_ctx` is used to store context during streaming. */
  otcrypto_sha2_context_t hash_ctx;
} hmac_extended_test_vector_t;

static hmac_extended_test_vector_t
    kHmacExtendedTestVectors[2 * ARRAYSIZE(kHmacTestVectors)];

/**
 * Populate `kHmacExtendedTestVectors` array with tests such that:
 * i) There are mixture of oneshot and streaming HMAC operations.
 * Each vector is run in 2 duplicates, one as a streaming and another
 * as a oneshot call.
 * ii) For streaming operations, messages are split into [1,4]
 * segments with randomly chosen indices.
 * iii) The order of tests is randomized.
 */
static void vectors_populate(void) {
  hmac_extended_test_vector_t *cur_ext_vec;
  hmac_test_vector_t *cur_vec;
  // First add streaming calls to the test array.
  for (size_t i = 0; i < ARRAYSIZE(kHmacTestVectors); i++) {
    cur_ext_vec = &kHmacExtendedTestVectors[i];
    cur_vec = &kHmacTestVectors[i];
    cur_ext_vec->hmac_test_vector = cur_vec;
    cur_ext_vec->original_idx = i;
    cur_ext_vec->segment_count = rand_testutils_gen32_range(1, 4);
    // Pick three indicies to break down message into 4 parts, but we only
    // use the first `segment_count` of them. The indices of a message
    // looks like below:
    // ix[0] = 0 ... idx[1] ... idx[2] ... idx[3] ...  message_len-1
    cur_ext_vec->segment_idx[0] = 0;
    cur_ext_vec->segment_idx[2] =
        rand_testutils_gen32_range(0, cur_vec->message.len);
    cur_ext_vec->segment_idx[1] =
        rand_testutils_gen32_range(0, cur_ext_vec->segment_idx[2]);
    cur_ext_vec->segment_idx[3] = rand_testutils_gen32_range(
        cur_ext_vec->segment_idx[2], cur_vec->message.len);
    cur_ext_vec->progress = kHmacTestNotStarted;
  }

  // Add oneshot calls of the same vectors.
  for (size_t i = 0; i < ARRAYSIZE(kHmacTestVectors); i++) {
    cur_ext_vec = &kHmacExtendedTestVectors[ARRAYSIZE(kHmacTestVectors) + i];
    cur_vec = &kHmacTestVectors[i];
    cur_ext_vec->hmac_test_vector = cur_vec;
    cur_ext_vec->original_idx = i;
    cur_ext_vec->segment_count = 0;
    cur_ext_vec->progress = kHmacTestNotStarted;
  }

  // Now we shuffle the order of vectors and print
  // how the new order maps to the original order from the test file.
  rand_testutils_shuffle(kHmacExtendedTestVectors,
                         sizeof(hmac_extended_test_vector_t),
                         ARRAYSIZE(kHmacExtendedTestVectors));
  for (size_t i = 0; i < ARRAYSIZE(kHmacExtendedTestVectors); i++) {
    cur_ext_vec = &kHmacExtendedTestVectors[i];
    cur_vec = cur_ext_vec->hmac_test_vector;
    LOG_INFO("Vector order = %d, original_vector_id = %d", i,
             cur_ext_vec->original_idx);
    LOG_INFO(
        "segment_count = %d, segment_idx = [%d, %d, %d, %d], message_len = %d",
        cur_ext_vec->segment_count, cur_ext_vec->segment_idx[0],
        cur_ext_vec->segment_idx[1], cur_ext_vec->segment_idx[2],
        cur_ext_vec->segment_idx[3],
        cur_ext_vec->hmac_test_vector->message.len);
  }
}

/**
 * Determines `hash_mode` for given SHA-2 test vectors.
 *
 * Note that for HMAC operations, mode information is part of the key struct,
 * hence this function is only used for hash vectors.
 *
 * @param test_vec The pointer to the test vector.
 * @param[out] hash_mode The determined hash_mode of the given test vector.
 */
static status_t get_hash_mode(hmac_test_vector_t *test_vec,
                              otcrypto_hash_mode_t *hash_mode) {
  switch (test_vec->test_operation) {
    case kHmacTestOperationSha256:
      *hash_mode = kOtcryptoHashModeSha256;
      return OK_STATUS();
    case kHmacTestOperationSha384:
      *hash_mode = kOtcryptoHashModeSha384;
      return OK_STATUS();
    case kHmacTestOperationSha512:
      *hash_mode = kOtcryptoHashModeSha512;
      return OK_STATUS();
    default:
      return INVALID_ARGUMENT();
  }
}

/**
 * Instantiate `hash_ctx` context object with hash/mac mode specified in
 * `current_test_vector`.
 *
 * @param hash_ctx Corresponding context for given `current_test_vector`.
 * @param current_test_vector Pointer to the hardcoded test vector.
 */
static status_t ctx_init(otcrypto_sha2_context_t *hash_ctx,
                         hmac_test_vector_t *current_test_vector) {
  // Populate `checksum` and `config.security_level` fields.
  current_test_vector->key.checksum =
      integrity_blinded_checksum(&current_test_vector->key);

  otcrypto_hash_mode_t hash_mode;

  switch (current_test_vector->test_operation) {
    case kHmacTestOperationSha256:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationSha384:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationSha512:
      TRY(get_hash_mode(current_test_vector, &hash_mode));
      TRY(otcrypto_sha2_init(hash_mode, hash_ctx));
      break;
    case kHmacTestOperationHmacSha256:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationHmacSha384:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationHmacSha512:
      TRY(otcrypto_hmac_init((otcrypto_hmac_context_t *)hash_ctx,
                             &current_test_vector->key));
      break;
    default:
      return INVALID_ARGUMENT();
  }
  return OK_STATUS();
}

/**
 * Run the test given by `current_test_vector` as oneshot operation.
 *
 * @param current_test_vector Pointer to the test vector.
 * @return The result of the operation.
 */
static status_t hmac_oneshot(hmac_test_vector_t *current_test_vector) {
  // Populate `checksum` and `config.security_level` fields.
  current_test_vector->key.checksum =
      integrity_blinded_checksum(&current_test_vector->key);

  // The test vectors already have the correct digest sizes hardcoded.
  size_t digest_len = current_test_vector->digest.len;
  // Allocate the buffer for the maximum digest size (which comes from SHA-512).
  uint32_t act_tag[512 / 32];
  otcrypto_hash_digest_t hash_digest = {
      .data = act_tag,
      .len = digest_len,
  };
  switch (current_test_vector->test_operation) {
    case kHmacTestOperationSha256:
      TRY(otcrypto_sha2_256(current_test_vector->message, &hash_digest));
      break;
    case kHmacTestOperationSha384:
      TRY(otcrypto_sha2_384(current_test_vector->message, &hash_digest));
      break;
    case kHmacTestOperationSha512:
      TRY(otcrypto_sha2_512(current_test_vector->message, &hash_digest));
      break;
    case kHmacTestOperationHmacSha256:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationHmacSha384:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationHmacSha512:
      TRY(otcrypto_hmac(
          &current_test_vector->key, current_test_vector->message,
          (otcrypto_word32_buf_t){.data = act_tag, .len = digest_len}));
      break;
    default:
      return INVALID_ARGUMENT();
  }
  LOG_INFO("Comparing result for %s.", current_test_vector->vector_identifier);
  CHECK_ARRAYS_EQ(act_tag, current_test_vector->digest.data, digest_len);
  return OK_STATUS();
}

/**
 * Feed messages to the streaming hash/hmac operation with context `hash_ctx`.
 * `current_test_vector` is used to determine message bytes to be used.
 *
 * This function feeds (approximately) the first or the second half of
 * `current_test_vector->message`.
 *
 * @param hash_ctx Corresponding context for given `current_test_vector`.
 * @param current_test_vector Pointer to the hardcoded test vector.
 * @param segment_start The starting index of the chosen segment.
 * @param segment_len The byte length of the chosen segment.
 * @return The result of the operation.
 */
static status_t feed_msg(otcrypto_sha2_context_t *hash_ctx,
                         hmac_test_vector_t *current_test_vector,
                         size_t segment_start, size_t segment_len) {
  otcrypto_const_byte_buf_t msg = {
      .data = &current_test_vector->message.data[segment_start],
      .len = segment_len,
  };

  switch (current_test_vector->test_operation) {
    case kHmacTestOperationSha256:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationSha384:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationSha512:
      TRY(otcrypto_sha2_update(hash_ctx, msg));
      break;
    case kHmacTestOperationHmacSha256:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationHmacSha384:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationHmacSha512:
      TRY(otcrypto_hmac_update((otcrypto_hmac_context_t *)hash_ctx, msg));
      break;
    default:
      return INVALID_ARGUMENT();
  }
  return OK_STATUS();
}

/**
 * Finalize the hash/hmac operation with context `hash_ctx`. Further necessary
 * parameters for the final call is reads from `current_test_vector`.
 *
 * @param hash_ctx Corresponding context for given `current_test_vector`.
 * @param current_test_vector Pointer to the hardcoded test vector.
 * @return The result of the operation.
 */
static status_t hmac_finalize(otcrypto_sha2_context_t *hash_ctx,
                              hmac_test_vector_t *current_test_vector) {
  // The test vectors already have the correct digest sizes hardcoded.
  size_t digest_len = current_test_vector->digest.len;
  // Allocate the buffer for the maximum digest size (which comes from SHA-512).
  uint32_t act_tag[512 / 32];
  otcrypto_word32_buf_t tag_buf = {
      .data = act_tag,
      .len = digest_len,
  };
  otcrypto_hash_digest_t hash_digest = {
      .data = act_tag,
      .len = digest_len,
  };
  switch (current_test_vector->test_operation) {
    case kHmacTestOperationSha256:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationSha384:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationSha512:
      TRY(otcrypto_sha2_final(hash_ctx, &hash_digest));
      break;
    case kHmacTestOperationHmacSha256:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationHmacSha384:
      OT_FALLTHROUGH_INTENDED;
    case kHmacTestOperationHmacSha512:
      TRY(otcrypto_hmac_final((otcrypto_hmac_context_t *)hash_ctx, tag_buf));
      break;
    default:
      return INVALID_ARGUMENT();
  }
  LOG_INFO("Comparing result for %s.", current_test_vector->vector_identifier);
  CHECK_ARRAYS_EQ(act_tag, current_test_vector->digest.data, digest_len);
  return OK_STATUS();
}

/**
 * Process a message segment for the given test vector.
 *
 * If the given test vector is already run and completed, this function
 * simply returns.
 *
 * If the given test vector is a oneshot call, then it is run and the
 * result is compared. The process of the vector is already update to done.
 *
 * If the given test vector is a starming call, then this function calls
 * the necessary init, update or final call. In the final call, the result
 * is compared with the expected result.
 *
 * @param test_ext_vec The test vector to run.
 * @return The result of the operation.
 *
 */
static status_t process_segment(hmac_extended_test_vector_t *test_ext_vec) {
  // If `test_ext_vec` is done, simply return.
  if (test_ext_vec->progress == kHmacTestDone) {
    return OK_STATUS();
  }

  // If `test_ext_vec` is one-shot, we need to call one-shot API.
  if (test_ext_vec->segment_count == 0 &&
      test_ext_vec->progress == kHmacTestNotStarted) {
    LOG_INFO("Invoking oneshot HMAC for vector #%d",
             test_ext_vec->original_idx);
    hmac_oneshot(test_ext_vec->hmac_test_vector);
    // Mark this test as complete
    test_ext_vec->progress = kHmacTestDone;
    return OK_STATUS();
  }

  // A sanity check: oneshot calls should not arrive here and `segment_count`
  // should be valid.
  if (test_ext_vec->segment_count == 0 || test_ext_vec->segment_count > 4) {
    return INVALID_ARGUMENT();
  }

  // If `test_ext_vec` is streaming, and this is the first call, then call init.
  if (test_ext_vec->progress == kHmacTestNotStarted) {
    LOG_INFO("Initializing HMAC stream for vector #%d",
             test_ext_vec->original_idx);
    ctx_init(&test_ext_vec->hash_ctx, test_ext_vec->hmac_test_vector);
    test_ext_vec->progress = kHmacTestFeedSegment1;
    return OK_STATUS();
  }

  if (test_ext_vec->progress == kHmacTestFinalize) {
    LOG_INFO("Finalizing HMAC stream for vector #%d",
             test_ext_vec->original_idx);
    hmac_finalize(&test_ext_vec->hash_ctx, test_ext_vec->hmac_test_vector);
    test_ext_vec->progress = kHmacTestDone;
    return OK_STATUS();
  }

  // Handle the last message specially
  if (test_ext_vec->progress == test_ext_vec->segment_count) {
    size_t segment_start =
        test_ext_vec->segment_idx[test_ext_vec->progress - 1];
    size_t segment_len =
        test_ext_vec->hmac_test_vector->message.len - segment_start;
    // Sanity check
    if (segment_start > test_ext_vec->hmac_test_vector->message.len) {
      return INVALID_ARGUMENT();
    }
    LOG_INFO("Streaming the last segment #%d of vector #%d.",
             test_ext_vec->progress, test_ext_vec->original_idx);
    feed_msg(&test_ext_vec->hash_ctx, test_ext_vec->hmac_test_vector,
             segment_start, segment_len);
    test_ext_vec->progress = kHmacTestFinalize;
    return OK_STATUS();
  }

  if (test_ext_vec->progress < test_ext_vec->segment_count) {
    size_t segment_start =
        test_ext_vec->segment_idx[test_ext_vec->progress - 1];
    size_t segment_len =
        test_ext_vec->segment_idx[test_ext_vec->progress] - segment_start;
    LOG_INFO("Streaming segment #%d of vector #%d.", test_ext_vec->progress,
             test_ext_vec->original_idx);
    feed_msg(&test_ext_vec->hash_ctx, test_ext_vec->hmac_test_vector,
             segment_start, segment_len);
    test_ext_vec->progress++;
    return OK_STATUS();
  }
  return INVALID_ARGUMENT();
}

/**
 * Run all vectors specified in `kHmacExtendedTestVectors` in parallel step
 * by step.
 */
static status_t run_test(void) {
  vectors_populate();

  for (size_t i = kHmacTestNotStarted; i < kHmacTestDone; i++) {
    for (size_t j = 0; j < ARRAYSIZE(kHmacExtendedTestVectors); j++) {
      TRY(process_segment(&kHmacExtendedTestVectors[j]));
    }
  }
  return OK_STATUS();
}

OTTF_DEFINE_TEST_CONFIG();
bool test_main(void) {
  LOG_INFO("Testing cryptolib SHA-2/HMAC with parallel multiple streams.");
  status_t test_result = OK_STATUS();
  EXECUTE_TEST(test_result, run_test);
  return status_ok(test_result);
}
